# 导入必要的库和模块
import numpy as np  # 用于科学计算
import matplotlib.pyplot as plt  # 用于绘制图形
from sklearn.datasets import load_digits  # 用于加载手写数字数据集
from sklearn.model_selection import train_test_split  # 用于分割训练集和测试集
from sklearn.neighbors import KNeighborsClassifier  # K近邻分类器
from sklearn.metrics import accuracy_score  # 用于计算准确率
import pickle  # 用于序列化和反序列化Python对象结构
from tqdm import tqdm  # 进度条库

# 加载数字数据集
digits = load_digits()
X = digits.data  # 特征数据
y = digits.target  # 标签数据

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
best_accuracy = 0  # 最佳准确率
best_k = 1  # 最佳k值
best_model = None  # 最佳knn模型

# 初始化一个列表以存储每个k值的准确率
accuracies = []  # 每个k值对应的准确率列表

# 尝试从1到40的k值，对于每个k值，训练knn模型，保存最佳准确率，k值和knn模型
for k in tqdm(range(1, 41), desc="Training KNN Models"):  # 使用tqdm创建进度条
    # 创建KNN模型实例
    knn = KNeighborsClassifier(n_neighbors=k)  # 使用当前k值初始化knn分类器
    
    # 训练模型
    knn.fit(X_train, y_train)  # 使用训练数据拟合模型
    
    # 在测试集上进行预测
    y_pred = knn.predict(X_test)  # 预测测试集标签
    
    # 计算准确率
    accuracy = accuracy_score(y_test, y_pred)  # 比较预测结果与真实标签，计算准确率
    
    # 更新最佳准确率、k值和模型
    if accuracy > best_accuracy:
        best_accuracy = accuracy  # 更新最高准确率
        best_k = k  # 更新最优k值
        best_model = knn  # 更新最优模型
    
    # 存储准确率
    accuracies.append(accuracy)  # 将当前k值下的准确率添加到列表中
    
    # 打印进度
    print(f"K={k}, Accuracy={accuracy:.4f}")  # 打印当前k值及对应的准确率

# 绘制准确率曲线
plt.figure(figsize=(10, 6))  # 设置绘图窗口大小
plt.plot(range(1, 41), accuracies, marker='o')  # 绘制准确率曲线
plt.title("Accuracy vs. Number of Neighbors (K)")  # 图表标题
plt.xlabel("Number of Neighbors K")  # x轴标签
plt.ylabel("Accuracy")  # y轴标签
plt.xticks(np.arange(1, 41, 2))  # 设置x轴刻度

# 绘制最佳K值的垂直线
plt.axvline(x=best_k, color='r', linestyle='--')  # 在最佳k值处画一条红色虚线
plt.text(best_k, accuracies[best_k-1], f'K={best_k}, Acc={best_accuracy:.4f}', fontsize=12, color='red')  # 添加文字注释

# 保存图像
plt.savefig('knn_accuracy_vs_k.pdf')  # 保存图表为PDF格式

# 将最佳KNN模型保存到二进制文件
with open(f'best_knn_model_k{best_k}.pkl', 'wb') as file:  # 使用pickle保存模型
    pickle.dump(best_model, file)

print(f"Best model saved to disk with K={best_k}")  # 打印模型保存信息

# 打印最佳准确率和相应的k值
print(f"\nBest K={best_k} with Accuracy={best_accuracy:.4f}")  # 打印最佳k值及其准确率